import os
os.environ["CUDA_VISIBLE_DEVICES"] = '2,3'
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
from modelscope import GenerationConfig
import jsonlines
import chromadb
from sentence_transformers import SentenceTransformer

question_json_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/question.json'
answer_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/submit_result.jsonl'
model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B'
content = []


client = chromadb.PersistentClient(path="utils/db")
collection = client.get_collection(name="my_collection")
embedding_model = SentenceTransformer('/datasets/fengjiahao/nlp/m3e-base/')

with jsonlines.open(question_json_path, "r") as json_file:
    for obj in json_file.iter(type=dict, skip_invalid=True):
        content.append(obj)


# Note: The default behavior now has injection attack prevention off.
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

# use bf16
# model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="cuda:0", trust_remote_code=True, bf16=True).eval()
# use cpu only
# model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="cpu", trust_remote_code=True).eval()
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
# 模型加载指定device_map='cuda:0'，更改成device_map='auto'会使用所有可用显卡

# Specify hyperparameters for generation
model.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)


def ask_llm(question):
    inputs = tokenizer('问题:%s\n答案:' % ( question), return_tensors='pt')
    inputs = inputs.to(model.device)
    pred = model.generate(**inputs)
    answer = tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)
    print(answer)
    return answer

for cont in content:
    question = cont['question']
    answer  = ask_llm(question)
    # print(question)

    cont['answer'] = answer.split('\n答案:')[1].replace('\n','')
    # 市盈率是最常用来评估股价水平是否合理的指标之一，是指股票价格与每股盈利的比率。...
with jsonlines.open(answer_path, "w") as json_file:
    json_file.write_all(content)